from typing import Dict, List, Tuple, Optional
from rainforeLearn.gomoku.v2.train.constants.train_stats_printer_constants import TrainingStatsPrinterConstants


class ThreatDetector:
    """威胁检测器 - 专门处理五子棋威胁分析"""
    
    def __init__(self, board_size: int):
        self.board_size = board_size
        self.directions = [(0, 1), (1, 0), (1, 1), (1, -1)]
        self.direction_names = ["horizontal", "vertical", "diagonal1", "diagonal2"]
    
    def analyze_threats(self, board, row: int, col: int, player: int) -> Dict:
        """分析指定位置的威胁情况"""
        opponent = -player
        threats_info = {
            "threats_detected": {},
            "multiple_threats": 0,
            "threat_penalty": 0
        }
        
        # 分析各方向威胁
        for i, (dr, dc) in enumerate(self.directions):
            threat_level = self._analyze_direction_threat(board, row, col, dr, dc, opponent)
            if threat_level:
                threats_info["threats_detected"][self.direction_names[i]] = threat_level
        
        # 计算多重威胁
        threats_info["multiple_threats"] = self._count_multiple_threats(board, row, col, opponent)
        
        # 计算威胁惩罚
        threats_info["threat_penalty"] = self._calculate_threat_penalty(threats_info)
        
        return threats_info
    
    def _analyze_direction_threat(self, board, r: int, c: int, dr: int, dc: int, opponent: int) -> Optional[str]:
        """分析特定方向的威胁等级"""
        positions = self._get_direction_positions(board, r, c, dr, dc)
        patterns = self._identify_threat_patterns(positions, opponent)
        return self._evaluate_threat_level(patterns)
    
    def _get_direction_positions(self, board, r: int, c: int, dr: int, dc: int) -> List[int]:
        """获取指定方向的位置序列"""
        positions = []
        for i in range(-3, 4):
            nr, nc = r + dr * i, c + dc * i
            if 0 <= nr < self.board_size and 0 <= nc < self.board_size:
                positions.append(board[nr, nc])
            else:
                positions.append(-999)  # 边界标记
        return positions
    
    def _identify_threat_patterns(self, positions: List[int], opponent: int) -> Dict[str, bool]:
        """识别威胁模式"""
        patterns = {
            'four_in_row': False,
            'open_four': False,
            'closed_four': False,
            'open_three': False,
            'closed_three': False,
            'open_two': False,
            'closed_two': False
        }
        
        consecutive_segments = self._find_consecutive_segments(positions, opponent)
        
        for segment in consecutive_segments:
            self._classify_segment_pattern(segment, patterns)
        
        return patterns
    
    def _find_consecutive_segments(self, positions: List[int], opponent: int) -> List[Dict]:
        """找到连续的棋子段"""
        segments = []
        current_count = 0
        start_idx = -1
        
        for i, pos in enumerate(positions):
            if pos == opponent:
                if current_count == 0:
                    start_idx = i
                current_count += 1
            else:
                if current_count > 0:
                    left_empty = (start_idx > 0 and positions[start_idx - 1] == 0)
                    right_empty = (i < len(positions) and positions[i] == 0)
                    
                    segments.append({
                        'length': current_count,
                        'start': start_idx,
                        'end': i - 1,
                        'left_empty': left_empty,
                        'right_empty': right_empty
                    })
                current_count = 0
        
        # 处理结尾的连续段
        if current_count > 0:
            left_empty = (start_idx > 0 and positions[start_idx - 1] == 0)
            segments.append({
                'length': current_count,
                'start': start_idx,
                'end': len(positions) - 1,
                'left_empty': left_empty,
                'right_empty': False
            })
        
        return segments
    
    def _classify_segment_pattern(self, segment: Dict, patterns: Dict[str, bool]):
        """根据连续段分类威胁模式"""
        length = segment['length']
        left_empty = segment['left_empty']
        right_empty = segment['right_empty']
        
        if length >= 4:
            patterns['four_in_row'] = True
        elif length == 3:
            if left_empty and right_empty:
                patterns['open_four'] = True
            elif left_empty or right_empty:
                patterns['closed_four'] = True
        elif length == 2:
            if left_empty and right_empty:
                patterns['open_three'] = True
            elif left_empty or right_empty:
                patterns['closed_three'] = True
        elif length == 1:
            if left_empty and right_empty:
                patterns['open_two'] = True
            elif left_empty or right_empty:
                patterns['closed_two'] = True
    
    def _evaluate_threat_level(self, patterns: Dict[str, bool]) -> Optional[str]:
        """评估威胁等级"""
        if patterns['four_in_row']:
            return TrainingStatsPrinterConstants.THREAT_CRITICAL
        elif patterns['open_four']:
            return TrainingStatsPrinterConstants.THREAT_HIGH
        elif patterns['closed_four'] or patterns['open_three']:
            return TrainingStatsPrinterConstants.THREAT_MEDIUM
        elif patterns['closed_three'] or patterns['open_two']:
            return TrainingStatsPrinterConstants.THREAT_LOW
        return None
    
    def _count_multiple_threats(self, board, r: int, c: int, opponent: int) -> int:
        """计算多重威胁数量"""
        threat_count = 0
        for dr, dc in self.directions:
            threat_level = self._analyze_direction_threat(board, r, c, dr, dc, opponent)
            if threat_level in [TrainingStatsPrinterConstants.THREAT_CRITICAL, 
                              TrainingStatsPrinterConstants.THREAT_HIGH, 
                              TrainingStatsPrinterConstants.THREAT_MEDIUM]:
                threat_count += 1
        return threat_count
    
    def _calculate_threat_penalty(self, threats_info: Dict) -> float:
        """计算威胁惩罚值"""
        penalty = 0.0
        threat_values = {
            TrainingStatsPrinterConstants.THREAT_CRITICAL: 2.0,
            TrainingStatsPrinterConstants.THREAT_HIGH: 1.0,
            TrainingStatsPrinterConstants.THREAT_MEDIUM: 0.5,
            TrainingStatsPrinterConstants.THREAT_LOW: 0.1
        }
        
        for threat_level in threats_info["threats_detected"].values():
            penalty += threat_values.get(threat_level, 0)
        
        # 多重威胁加重惩罚
        if threats_info["multiple_threats"] > 1:
            penalty *= (1 + threats_info["multiple_threats"] * 0.5)
        
        return penalty